"""
RAG 文档处理模块
使用 LangChain 实现文档加载和文本切分功能
支持 PDF、TXT、DOCX 文件类型
"""
import os
from typing import List, Optional, Dict, Any
from pathlib import Path

from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    Docx2txtLoader,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document


class RAGDocumentHandler:
    """
    RAG 文档处理器
    
    功能：
    1. 根据文件类型自动选择合适的 LangChain Loader
    2. 加载文档并转换为 Document 对象
    3. 对文档进行智能切分，支持重叠和自定义参数
    4. 提供统一的接口处理多种文件格式
    """
    
    # 支持的文件类型映射
    SUPPORTED_EXTENSIONS = {
        '.pdf': 'pdf',
        '.txt': 'txt',
        '.docx': 'docx',
    }
    
    def __init__(
        self,
        chunk_size: int = 1000,
        chunk_overlap: int = 100,
        separators: Optional[List[str]] = None,
        encoding: str = 'utf-8'
    ):
        """
        初始化 RAG 文档处理器
        
        :param chunk_size: 每个文本块的最大字符数，默认 1000
        :param chunk_overlap: 文本块之间的重叠字符数，默认 100
        :param separators: 自定义分隔符列表，用于文本切分
                         如果为 None，使用 RecursiveCharacterTextSplitter 的默认分隔符
        :param encoding: 文本文件的编码，默认 'utf-8'
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.encoding = encoding
        
        # 初始化文本切分器
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            length_function=len,
        )
    
    def load_document(self, file_path: str) -> List[Document]:
        """
        加载文档并返回 Document 对象列表
        
        :param file_path: 文件路径
        :return: Document 对象列表
        :raises FileNotFoundError: 如果文件不存在
        :raises ValueError: 如果文件类型不支持
        :raises RuntimeError: 如果加载过程中发生错误
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件未找到: {file_path}")
        
        file_path_obj = Path(file_path)
        file_extension = file_path_obj.suffix.lower()
        
        if file_extension not in self.SUPPORTED_EXTENSIONS:
            raise ValueError(
                f"不支持的文件类型: {file_extension}。"
                f"支持的类型: {', '.join(self.SUPPORTED_EXTENSIONS.keys())}"
            )
        
        try:
            loader = self._get_loader(file_path, file_extension)
            documents = loader.load()
            return documents
        except Exception as e:
            raise RuntimeError(f"加载文档失败: {file_path}。错误: {e}") from e
    
    def _get_loader(self, file_path: str, file_extension: str):
        """
        根据文件扩展名获取对应的 LangChain Loader
        
        :param file_path: 文件路径
        :param file_extension: 文件扩展名（包含点号，如 '.pdf'）
        :return: LangChain Loader 实例
        :raises ValueError: 如果文件类型不支持
        """
        if file_extension not in self.SUPPORTED_EXTENSIONS:
            raise ValueError(f"不支持的文件类型: {file_extension}")
        
        file_type = self.SUPPORTED_EXTENSIONS[file_extension]
        
        if file_type == 'pdf':
            return PyPDFLoader(file_path)
        elif file_type == 'txt':
            return TextLoader(file_path, encoding=self.encoding)
        elif file_type == 'docx':
            return Docx2txtLoader(file_path)
        else:
            raise ValueError(f"未实现的文件类型: {file_type}")
    
    def split_documents(self, documents: List[Document]) -> List[Document]:
        """
        将文档列表切分成更小的文本块
        
        :param documents: Document 对象列表
        :return: 切分后的 Document 对象列表
        """
        if not documents:
            return []
        
        try:
            split_docs = self.text_splitter.split_documents(documents)
            return split_docs
        except Exception as e:
            raise RuntimeError(f"文档切分失败。错误: {e}") from e
    
    def load_and_split(self, file_path: str) -> List[Document]:
        """
        加载文档并自动切分（便捷方法）
        
        :param file_path: 文件路径
        :return: 切分后的 Document 对象列表
        """
        documents = self.load_document(file_path)
        split_documents = self.split_documents(documents)
        return split_documents
    
    def get_document_texts(self, documents: List[Document]) -> List[str]:
        """
        从 Document 对象列表中提取纯文本内容
        
        :param documents: Document 对象列表
        :return: 文本字符串列表
        """
        return [doc.page_content for doc in documents]
    
    def get_document_metadata(self, documents: List[Document]) -> List[Dict[str, Any]]:
        """
        从 Document 对象列表中提取元数据
        
        :param documents: Document 对象列表
        :return: 元数据字典列表
        """
        return [doc.metadata for doc in documents]
    
    def process_file(
        self,
        file_path: str,
        return_texts: bool = False,
        return_metadata: bool = False
    ) -> Dict[str, Any]:
        """
        处理文件的完整流程：加载 -> 切分 -> 返回结果
        
        :param file_path: 文件路径
        :param return_texts: 是否返回纯文本列表
        :param return_metadata: 是否返回元数据列表
        :return: 包含处理结果的字典
        """
        split_documents = self.load_and_split(file_path)
        
        result = {
            'documents': split_documents,
            'chunk_count': len(split_documents),
        }
        
        if return_texts:
            result['texts'] = self.get_document_texts(split_documents)
        
        if return_metadata:
            result['metadata'] = self.get_document_metadata(split_documents)
        
        return result


class DocumentLoaderFactory:
    """
    文档加载器工厂类
    提供静态方法创建不同类型的文档加载器
    """
    
    @staticmethod
    def create_pdf_loader(file_path: str) -> PyPDFLoader:
        """
        创建 PDF 加载器
        
        :param file_path: PDF 文件路径
        :return: PyPDFLoader 实例
        """
        return PyPDFLoader(file_path)
    
    @staticmethod
    def create_txt_loader(file_path: str, encoding: str = 'utf-8') -> TextLoader:
        """
        创建文本文件加载器
        
        :param file_path: 文本文件路径
        :param encoding: 文件编码，默认 'utf-8'
        :return: TextLoader 实例
        """
        return TextLoader(file_path, encoding=encoding)
    
    @staticmethod
    def create_docx_loader(file_path: str) -> Docx2txtLoader:
        """
        创建 DOCX 加载器
        
        :param file_path: DOCX 文件路径
        :return: Docx2txtLoader 实例
        """
        return Docx2txtLoader(file_path)
    
    @staticmethod
    def create_loader_by_extension(
        file_path: str,
        encoding: str = 'utf-8'
    ):
        """
        根据文件扩展名自动创建对应的加载器
        
        :param file_path: 文件路径
        :param encoding: 文本文件编码，默认 'utf-8'
        :return: 对应的 LangChain Loader 实例
        :raises ValueError: 如果文件类型不支持
        """
        file_path_obj = Path(file_path)
        file_extension = file_path_obj.suffix.lower()
        
        if file_extension == '.pdf':
            return DocumentLoaderFactory.create_pdf_loader(file_path)
        elif file_extension == '.txt':
            return DocumentLoaderFactory.create_txt_loader(file_path, encoding)
        elif file_extension == '.docx':
            return DocumentLoaderFactory.create_docx_loader(file_path)
        else:
            raise ValueError(
                f"不支持的文件类型: {file_extension}。"
                f"支持的类型: .pdf, .txt, .docx"
            )


class TextSplitterConfig:
    """
    文本切分器配置类
    提供预定义的切分器配置
    """
    
    @staticmethod
    def create_default_splitter(
        chunk_size: int = 1000,
        chunk_overlap: int = 100
    ) -> RecursiveCharacterTextSplitter:
        """
        创建默认的递归字符文本切分器
        
        :param chunk_size: 块大小
        :param chunk_overlap: 重叠大小
        :return: RecursiveCharacterTextSplitter 实例
        """
        return RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
        )
    
    @staticmethod
    def create_markdown_splitter(
        chunk_size: int = 1000,
        chunk_overlap: int = 100
    ) -> RecursiveCharacterTextSplitter:
        """
        创建适用于 Markdown 的文本切分器
        
        :param chunk_size: 块大小
        :param chunk_overlap: 重叠大小
        :return: RecursiveCharacterTextSplitter 实例（使用 Markdown 分隔符）
        """
        separators = [
            "\n\n## ",
            "\n\n### ",
            "\n\n",
            "\n",
            " ",
            "",
        ]
        return RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            length_function=len,
        )
    
    @staticmethod
    def create_code_splitter(
        chunk_size: int = 1000,
        chunk_overlap: int = 100
    ) -> RecursiveCharacterTextSplitter:
        """
        创建适用于代码的文本切分器
        
        :param chunk_size: 块大小
        :param chunk_overlap: 重叠大小
        :return: RecursiveCharacterTextSplitter 实例（使用代码分隔符）
        """
        separators = [
            "\n\nclass ",
            "\n\ndef ",
            "\n\n    def ",
            "\n\n",
            "\n",
            " ",
            "",
        ]
        return RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=separators,
            length_function=len,
        )

